/************************************************************************/
/*                                                                      */
/*  Typecast Project                                                    */
/*                                                                      */
/*  User mode routing daemon                                            */
/*  Author: Eusden Shing, UCLA                                          */
/*                                                                      */
/************************************************************************/

#include<sys/stat.h>
#include<unistd.h>
#include<stdio.h>
#include<stdlib.h>
#include<assert.h>
#include<sys/socket.h>
#include<sys/types.h>
#include<string.h>
#include<asm/types.h>
#include<linux/netlink.h>
#include<linux/socket.h>
#include<arpa/inet.h>
#include<netinet/in.h>

#include <linux/netfilter.h>
#include <linux/ip.h>
#include <linux/udp.h>

#include "libipq.h"

#include "bloom_filter.h"
#include "type_routingtable.h"
#include "type_announce.h"
#include "typecast_config.h"
#include "LBF_update.h"
#include "ipq_process.h"

#include <time.h>
#define MAX_PAYLOAD 1024
#define MAX_BUFFER_SIZE 1024
#define MAX_NAME 255
#define TYPECAST_TA_PORT 2003
#define TYPECAST_LBF_PORT 2004


#define TA_REFRESH_PERIOD 5 


#define max(i, j)  ( i > j ? i : j)

//////////////////////////////////////////////////////////////////////////
//  Global Variables
//////////////////////////////////////////////////////////////////////////

const int MULTICAST_TTL = 64;

multicast_group_node * multicast_group_list_head = NULL;
multicast_group_node * multicast_group_list_tail = NULL;
int multicast_group_list_count = 0;

int receive_sock_fd = 0; 
int update_sock_fd = 0; 
int send_sock_fd = 0; 

struct ipq_handle * pipqHandle;

//////////////////////////////////////////////////////////////////////////
// ReceiveTypeAnnounce
//   Spawns a new thread that receives Type-Announce packet from all 
//   neighbors from all multicast groups (?).
//   Returns 0 on success.  Returns nonzero on failure.
int ReceiveTypeAnnounce();

//////////////////////////////////////////////////////////////////////////
// SendTypeAnnounce
//   Sends a Type-Announce packet to all neighbors on all multicast groups.
//   Returns 0 on success.  Returns nonzero on failure.
int SendTypeAnnounce();

//////////////////////////////////////////////////////////////////////////
// ConstructReceiveSockets
//   Initializes the sockets we need for receiving.
//   Returns 0 on success.  Returns nonzero on failure.
int ConstructReceiveSockets();

//////////////////////////////////////////////////////////////////////////
// ConstructUpdateSockets
//   Initializes the sockets we need for updating the LBF.
//   Returns 0 on success.  Returns nonzero on failure.
int ConstructUpdateSockets();

//////////////////////////////////////////////////////////////////////////
// ConstructSendSockets
//   Initializes the sockets we need for sending.
//   Returns 0 on success.  Returns nonzero on failure.
int ConstructSendSockets();

//////////////////////////////////////////////////////////////////////////
// ConstructIPQHandler
//   Returns 0 on success.  Returns nonzero on failure.
int ConstructIPQHandler();

//////////////////////////////////////////////////////////////////////////
// ReceiveIPQPacket
//   Returns 0 on success.  Returns nonzero on failure.
int ReceiveIPQPacket();

//////////////////////////////////////////////////////////////////////////
// CleanUp
void CleanUp();


/* Simple function (based on R. Stevens) to calculate UDP header checksum */
static inline u_int16_t in_udp_csum(unsigned long saddr,
                                    unsigned long daddr,
                                    unsigned short len,
                                    unsigned short proto,
                                    unsigned short *buf)
{
    unsigned short *w;
    int nleft = ntohs(len);
    int sum = 0;
    struct {
        unsigned short zero:8;
        unsigned short proto:8;
    } proto2 = {0, proto};

    w = (unsigned short *)&saddr;
    sum += *w++;
    sum += *w;

    w = (unsigned short *)&daddr;
    sum += *w++;
    sum += *w;

    sum += *(unsigned short *)&proto2;
    sum += len;

    w = buf;

    while(nleft > 1) {
        sum += *w++;
        nleft -= 2;
    }

    if (nleft == 1)
        sum += *(unsigned char*)w;

    sum = (sum >> 16) + (sum & 0xffff);
    sum += (sum >> 16);          

    return ~sum;

}



//////////////////////////////////////////////////////////////////////////
//////////////////////////////////////////////////////////////////////////
int main(int argc, char **argv)
{
    struct timeval tv; 
    fd_set rfds; 
    fd_set readers;         
    int nfds; 
    
    int iRet; 
                
    printf("Starting...\n");

    // Load routing information
	read_typecast_config_file(); 
    printf("ProcessDataFile succeeded.\n");	    
    local_ip_addr = get_local_ip_address();
	
    // Set up sockets
    if (ConstructSendSockets() != 0)
    {
        printf("Initialization of send sockets failed.\n");
        CleanUp();
        return -1;
    }
    printf("ConstructSendSockets succeeded.\n");

    if (ConstructReceiveSockets() != 0)
    {
        printf("Initialization of receive sockets failed.\n");
        CleanUp();
        return -1;
    }
    printf("ConstructReceiveSockets succeeded.\n");

    if (ConstructUpdateSockets() != 0)
    {
        printf("Initialization of update sockets failed.\n");
        CleanUp();
        return -1;
    }
    printf("ConstructUpdateSockets succeeded.\n");


    if (ConstructIPQHandler() != 0)
    {
        printf("Initialization of IPQ handler failed.\n");
        CleanUp();
        return -1;
    }
    printf("ConstructIPQHandler succeeded.\n");

    // Enter loop for periodic sending and receiving of Type-Announce, and receiving of IPQ packets
    tv.tv_sec = TA_REFRESH_PERIOD; 
    tv.tv_usec = 0;
    FD_ZERO(&readers); 
    FD_SET(receive_sock_fd, &readers);
    FD_SET(update_sock_fd, &readers);
    FD_SET(pipqHandle->fd, &readers);
    nfds = max(update_sock_fd, max(receive_sock_fd, pipqHandle->fd)) + 1;         
    
    for (;;)
    {
        memcpy((char *) &rfds, (char *) &readers, sizeof(rfds));        

        // See if there is anything to read
        if ((iRet = select(nfds, &rfds, NULL, NULL, &tv)) > 0)          
        {
            if (FD_ISSET(receive_sock_fd, &rfds))
            {
                if (ReceiveTypeAnnounce() < 0)
                {
                    printf("ReceiveTypeAnnounce failed.\n");
                    CleanUp();
                    return -1;
                }
                printf("ReceiveTypeAnnounce succeeded.\n");
            }
            if (FD_ISSET(update_sock_fd, &rfds))
            {
                if (UpdateLBF() < 0)
                {
                    printf("UpdateLBF failed.\n");
                    CleanUp();
                    return -1;
                }
                printf("UpdateLBF succeeded.\n");
            }
            if (FD_ISSET(pipqHandle->fd, &rfds))
            {
                if (ReceiveIPQPacket() < 0)
                {
                    printf("ReceiveIPQPacket failed.\n");
                    CleanUp();
                    return -1;
                }
                printf("ReceiveIPQPacket succeeded.\n");
            }
            // Relying on Linux Kernel to update tv in the select() call. 
        } else { 
            if (iRet == 0) { // timeout
                if (SendTypeAnnounce() != 0)
                {
                    printf("SendTypeAnnounce failed.\n");
                    CleanUp();
                    return -1;
                }
                printf("SendTypeAnnounce succeeded.\n");
                tv.tv_sec = TA_REFRESH_PERIOD; 
                tv.tv_usec = 0;                 
            } else { 
                perror("Error in select().\n");
                return -1;
            }
        }

        printf("--------------------------------\n\n");
        DumpTypeCastRoutingTable(stdout);
        printf("--------------------------------\n\n");
    }

    return 0;
}

void CleanUp()
{
    if (pipqHandle != NULL)
    {
        ipq_destroy_handle(pipqHandle);
        pipqHandle = NULL;
    }

    if (receive_sock_fd != 0)
    {
        close(receive_sock_fd);
        receive_sock_fd = 0;
    }

    if (update_sock_fd != 0)
    {
        close(update_sock_fd);
        update_sock_fd = 0;
    }

    if (send_sock_fd != 0)
    {
        close(send_sock_fd);
        send_sock_fd = 0;
    }
}

int ConstructReceiveSockets()
{
    // Initialize socket
    receive_sock_fd = socket(PF_INET, SOCK_RAW, IPPROTO_UDP); 
    if (receive_sock_fd < 0)
    { 
        perror("Error creating receiving socket.\n"); 
        return -1; 
    }

    const int on = 1; 
    if (setsockopt(receive_sock_fd, SOL_SOCKET, SO_BROADCAST, &on, sizeof(on)) < 0)
    {
        perror("Error when setsockopt in ConstructReceiveSockets\n");
        return -1;
    }

    if (setsockopt(receive_sock_fd, SOL_IP, IP_HDRINCL, &on, sizeof(on)) < 0)
    {
        perror("Error when setsockopt IP_HDRINCL in ConstructReceiveSockets\n");
        return -1;
    }


    // Bind the server.
    struct sockaddr_in serverName = { 0 };
    serverName.sin_family = AF_INET;
    serverName.sin_addr.s_addr = htonl(INADDR_ANY);
    serverName.sin_port = htons(TYPECAST_TA_PORT);

    if (bind(receive_sock_fd, (struct sockaddr *) &serverName, sizeof(serverName)) < 0)
    { 
        perror("bind()\n");
        return -1; 
    }
    
    return 0;
}

int ConstructUpdateSockets()
{
    // Initialize socket
    update_sock_fd = socket(PF_INET, SOCK_DGRAM, 0); 
    if (update_sock_fd < 0)
    { 
        perror("Error creating updating socket.\n"); 
        return -1; 
    }

    // Bind the server.
    struct sockaddr_in serverName = { 0 };
    serverName.sin_family = AF_INET;
    serverName.sin_addr.s_addr = htonl(INADDR_ANY);
    serverName.sin_port = htons(TYPECAST_LBF_PORT);

    if (bind(update_sock_fd, (struct sockaddr *) &serverName, sizeof(serverName)) < 0)
    { 
        perror("bind() update_sock_fd\n");
        return -1; 
    }

    return 0;
}

int ConstructSendSockets()
{
    // Initialize socket
    send_sock_fd = socket(PF_INET, SOCK_RAW, IPPROTO_UDP); 
    if (send_sock_fd < 0)
    { 
        perror("Error creating updating socket.\n"); 
        return -1; 
    }

    int on = 1;
    if (setsockopt(send_sock_fd, SOL_SOCKET, SO_BROADCAST, &on, sizeof(on)) < 0)
    {
        perror("Error in setsockopt in ConstructSendSockets.\n"); 
        return -1; 
    }

    if (setsockopt(send_sock_fd, SOL_IP, IP_HDRINCL, &on, sizeof(on)) < 0)
    {
        perror("Error when setsockopt in ConstructReceiveSockets\n");
        return -1;
    }

    return 0;
}

int ReceiveTypeAnnounce()
{
    int msglen = 0;
    struct sockaddr_in src_addr = { 0 };
    socklen_t addrlen = sizeof(src_addr);
    char buf[MAX_BUFFER_SIZE];

    memset(buf, 0, MAX_BUFFER_SIZE);

    msglen = recvfrom(receive_sock_fd, buf, sizeof(buf), 0, (struct sockaddr *) &src_addr, (socklen_t *) &addrlen);
    printf("    recvfrom ReceiveTypeAnnounce msglen = %d.\n", msglen);
    if (msglen < 0) 
    { 
        perror("Error when recvfrom.\n");
        return -1;
    }

    u_int32_t srcAddr = ntohl(src_addr.sin_addr.s_addr);
    if (srcAddr != local_ip_addr)
    {
        // Process the packet.
        ProcessTypeAnnouncePacket(buf + sizeof(struct iphdr) + sizeof(struct udphdr), srcAddr);
    }

    return 0;
}

int UpdateLBF()
{
    int msglen = 0;
    struct sockaddr_in src_addr = { 0 };
    socklen_t addrlen = sizeof(src_addr);
    char buf[MAX_BUFFER_SIZE];

    memset(buf, 0, MAX_BUFFER_SIZE);

    msglen = recvfrom(update_sock_fd, buf, sizeof(buf), 0, (struct sockaddr *) &src_addr, (socklen_t *) &addrlen);
    if (msglen < 0) 
    { 
        perror("Error when recvfrom.\n");
        return -1;
    }

    ProcessLBFUpdate(buf);

    return 0;
}

int SendTypeAnnounce()
{
    char buf[MAX_BUFFER_SIZE]; 
    memset(buf, 0, MAX_BUFFER_SIZE);

    static int first_time = 1;
    // Initialize seqno to a random number on the first packet.
    if (first_time == 1)
    {
        first_time = 0; 
        seqno = random(); 
    }

    // Aggregate all the ABFs.
    UpdateABFs();

    // Fill the buffer with header and routing information.
    int packet_size = GenerateTypeAnnouncePacket(buf);

    // Fill in the entire packet with header now.
    char packet[MAX_BUFFER_SIZE];
    memset(packet, 0, MAX_BUFFER_SIZE);

    struct iphdr *ip;
    struct udphdr *udp;
    int totlen = 0;

    totlen = sizeof(struct iphdr) + sizeof(struct udphdr) + packet_size;

    ip = (struct iphdr *) packet;
    ip->version = IPVERSION;
    ip->ihl = sizeof(struct iphdr) >> 2;
    ip->tos = IPTOS_PREC_NETCONTROL | IPTOS_LOWDELAY; /* Highest priority */
    ip->tot_len = htons(totlen);
    ip->id = 0; /* Let IP set */
    ip->frag_off = htons(0x4000); /* Set Do Not Fragment bit */
    ip->ttl = 16;    // !!!! TODO !!!! What do I set this to?
    ip->protocol = IPPROTO_UDP;
    ip->saddr = htonl(get_local_ip_address());  // !!!!! TODO !!!! Is this correct?
    ip->daddr = htonl(inet_addr("255.255.255.255"));

    ip->check = 0; /* Let kernel calculate */

    udp = (struct udphdr *) (packet + sizeof(struct iphdr));

    udp->source = htons(TYPECAST_TA_PORT);
    udp->dest = htons(TYPECAST_TA_PORT);
    udp->len = htons(packet_size + sizeof(struct udphdr));

    udp->check = in_udp_csum(ip->saddr, ip->daddr, udp->len, 
        ip->protocol, (unsigned short *)udp);

    memcpy(packet + sizeof(struct iphdr) + sizeof(struct udphdr),
           buf, packet_size);

    // Set up broadcast destination
    struct sockaddr_in dst_addr; 
    memset(&dst_addr, 0, sizeof(dst_addr));
    dst_addr.sin_family = AF_INET; 
    dst_addr.sin_addr.s_addr = htonl(inet_addr("255.255.255.255")); 
    dst_addr.sin_port = htons(TYPECAST_TA_PORT); 


    ssize_t msglen = 0; 
    msglen = sendto(send_sock_fd, packet, totlen, 0, (struct sockaddr *) &dst_addr, sizeof(dst_addr) );
    if (msglen < 0) 
    { 
        perror("Error when sendto.\n"); 
        return -1; // Too bad if it fails.
    }

    printf("sendto: sent with msglen = %d.\n", msglen);

    return 0;
}

int ConstructIPQHandler()
{
    int ipqRet;

    pipqHandle = ipq_create_handle(0, PF_INET);
    if (pipqHandle == NULL)
    {
        ipq_perror("Failed ipq_create_handle:");
        return -1;
    }

    ipqRet = ipq_set_mode(pipqHandle, IPQ_COPY_PACKET, MAX_BUFFER_SIZE);
    if (ipqRet < 0)
    {
        ipq_perror("Failed ipq_set_mode:");
        return -1;
    }

    return 0;
}

int ReceiveIPQPacket()
{
    return ProcessIPQPacket(pipqHandle);
}
